#include "generator.h"

const string Generator::stackingsymbol = "Stacking";
const string Generator::singlelayersymbol = "SingleLayer";
const string Generator::gongalsymbol = "GonGAL";
const string Generator::paddingsymbol = "Padding";
const string Generator::randomholecenterradiusssymbol = "RandomHoleCenterRadius";
const string Generator::recursiveedgeetchssymbol = "RecursiveEdgeEtch";
const string Generator::inftriangulargalsymbol = "InfTriangularGAL";
const string Generator::inftriangulargalrectanglesymbol = "InfTriangularGALRectangle";
const string Generator::triangulargalrectanglesymbol = "TriangularGALRectangle";

Generator::Generator(const symbol &str, const variant &param, const rule::list &opts) 
	: stacking(singlelayer), 
	padding(automatic), 
	userpadding(0), 
	randomholes(false), 
	randomradius(0.),
	recursiveedgeetch(0)
{
	const size_t rulesize = opts.size();
	for(size_t r = 0; r < rulesize; r++)
	{
		if(!opts[r].first.isSymbolType())
			continue;
		
		symbol sym = opts[r].first.toSymbol();
		variant var = opts[r].second;
		
		if(sym.name() == stackingsymbol && var.isSymbolType())
		{
			if(var.toSymbol().name() == singlelayersymbol)
				stacking = singlelayer;
			else if(var.toSymbol().name() == gongalsymbol)
				stacking = gongal;
		}
		else if(sym.name() == paddingsymbol)
		{
			if(var.isSymbolType() && var.toSymbol().getType() == symbol::predefined)
			{
				if(var.toSymbol().reference() == symbol::automatic)
					padding = automatic;
				else if(var.toSymbol().reference() == symbol::none)
					padding = none;
			}
			else if(var.isNumericType())
			{
				padding = userdefined;
				userpadding = var.toLong();
			}
		}
		else if(sym.name() == randomholecenterradiusssymbol)
		{
			if(var.isSymbolType() && var.toSymbol().getType() == symbol::predefined)
			{
				if(var.toSymbol().reference() == symbol::automatic)
				{
					randomholes = true;
					randomradius = 1.;
				}
				else if(var.toSymbol().reference() == symbol::none)
					randomholes = false;
			}
			else if(var.isNumericType())
			{
				randomradius = var.toDouble();

				randomholes = true;
			}	
		}
		else if(sym.name() == recursiveedgeetchssymbol)
		{
			if(var.isSymbolType() && var.toSymbol().getType() == symbol::predefined && var.toSymbol().reference() == symbol::none)
				recursiveedgeetch = 0;
			else if(var.isNumericType() && var.toLong() > 0)
				recursiveedgeetch = var.toLong();
		}
	}
	
	if(str.name() == inftriangulargalsymbol && param.isNumericListType())
	{
		if(param.toNumericList().size() < 2)
			return;

		size_t L = param.toNumericList()[0];
		double R = param.toNumericList()[1];
		
		atomsvec = inftriangulargal_atoms(L, R, stacking);
		pvec = inftriangulargal_periodicity(L, R, stacking);
	}
	else if(str.name() == inftriangulargalrectanglesymbol && param.isNumericListType())
	{
		if(param.toNumericList().size() < 2)
			return;

		size_t L = param.toNumericList()[0];
		double R = param.toNumericList()[1];
		size_t Nx = 1;
		size_t Ny = 1;
		if(param.toNumericList().size() >= 4)
		{
			Nx = param.toNumericList()[2];
			Ny = param.toNumericList()[3];
		}

		if(randomholes)
			atomsvec = inftriangulargalrectangle_atoms_randomholes(L, R, Nx, Ny, stacking, randomradius);
		else
			atomsvec = inftriangulargalrectangle_atoms(L, R, Nx, Ny, stacking);

		pvec = inftriangulargalrectangle_periodicity(L, R, Nx, Ny, stacking);
	}
	else if(str.name() == triangulargalrectanglesymbol && param.isNumericListType())
	{
		if(param.toNumericList().size() < 2)
			return;

		size_t L = param.toNumericList()[0];
		double R = param.toNumericList()[1];
		size_t Nx = 1;
		size_t Ny = 1;
		if(param.toNumericList().size() >= 4)
		{
			Nx = param.toNumericList()[2];
			Ny = param.toNumericList()[3];
		}

		if(padding != automatic && userpadding == 0)
			atomsvec = triangulargalrectangle_atoms(L, R, Nx, Ny, stacking);
		else if(padding == automatic)
			atomsvec = triangulargalrectangle_atoms(L, R, Nx, Ny, stacking,(size_t)R+1);
	}

	while(recursiveedgeetch > 0)
	{
		recursive_edge_etching(atomsvec,pvec,recursiveedgeetch);
		recursiveedgeetch--;
	}
}
	
const std::vector<arma::ivec4> & Generator::atoms() { return atomsvec; }
const std::vector<arma::ivec4> & Generator::periodicity() { return pvec; }

#include "atom.h"	

std::vector<arma::ivec4> Generator::inftriangulargal_atoms(size_t L, double R, const stackingtype &st)
{
	// First generate the hexagon of atoms:

	std::vector<arma::ivec4> atoms;
	atoms.reserve(4*L*L*4);

	static const auto atomGen = [] (size_t i,size_t j,size_t h,size_t l) 
	{
		arma::ivec4 result;
		result(0) = i;
		result(1) = j;
		result(2) = h;
		result(3) = l;
		return result;
	};
	
	for(size_t l = 0; l < 2; l++)
	{
		for(size_t j = 0; j < 2*L; j++)
		{
			for(size_t i = std::max((long)0,(long)j-(long)L+( (l == 1) ? 0 : 1 )); i < std::min(2*(long)L,(long)j+(long)L+( (l == 1) ? 0 : 1 )); i++)
			{
				for(size_t h = 0; h < ((st != singlelayer) ? 2: 1); h++)
				{
					atoms.push_back(atomGen(i,j,h,l));
				}
			}
		} 
	}

	// Now we make the hole.
	arma::vec2 holecent;
	holecent(0) = (double)L-0.5;
	holecent(1) = (double)0.5/sqrt(3.);
	
	std::vector<size_t> badatoms;
	size_t size = atoms.size();
	badatoms.reserve(size);
	
	for(size_t a = 0; a < size; a++)
	{
		if(arma::norm(holecent-Atom::changeBasis(atoms[a])(arma::span(0,1)),2) < R && ( st == singlelayer || atoms[a](2) == 0))
			badatoms.push_back(a);
	}

	Atom::filter::erase_indices(atoms,badatoms);

	return atoms;
}

std::vector<arma::ivec4> Generator::inftriangulargal_periodicity(size_t L, double R, const stackingtype &st)
{
	std::vector<arma::ivec4> per;
	
	static const auto perGen = [] (size_t i,size_t j) 
	{
		arma::ivec4 result;
		result(0) = i;
		result(1) = j;
		result(2) = 0;
		result(3) = 0;
		return result;
	};

	per.push_back(perGen(2*(long)L,(long)L));
	per.push_back(perGen((long)L,2*(long)L));
	return per;
}

std::vector<arma::ivec4> Generator::singletriangulargalrectangle_atoms(size_t L, double R, const stackingtype &st, arma::uvec padd)
{
	// First generate the rectangle of atoms:

	std::vector<arma::ivec4> atoms;
	atoms.reserve((3*L+padd(0)+padd(1))*(2*L+padd(2)+padd(3))*4);

	static const auto atomGen = [] (long i,long j,long h,long l) 
	{
		arma::ivec4 result;
		result(0) = i;
		result(1) = j;
		result(2) = h;
		result(3) = l;
		return result;
	};
	
	for(long l = 0; l < 2; l++)
	{
		for(long j = -(long)L-2*(long)padd(0)-2*(long)padd(1)-2*(long)padd(2)-2*(long)padd(3); j < 8*(long)L; j++)
		{
			for(long i = std::max(-j,j-( (l == 1) ? 1 : 0 )); i < std::min(2*(long)L+j-1+( (l == 1) ? 0 : 1 ) + 2*(long)padd(2) + 2*(long)padd(3),6*(long)L-j + 2*(long)padd(0) + 2*(long)padd(1)); i++)
			{
				for(long h = 0; h < ((st != singlelayer) ? 2: 1); h++)
				{
					atoms.push_back(atomGen(i,j,h,l));
				}
			}
		} 
	}

	// Now we make the holes.
	arma::vec2 holecent;
	holecent(0) = -((long)L % 2)*0.5+(long)padd(0);
	holecent(1) = 0.5/sqrt(3.)+((long)L-1+2*(long)padd(2))*0.5*sqrt(3.);
	
	std::vector<arma::ivec4> per = inftriangulargal_periodicity(L, R, st);

	arma::vec2 holeper1 = Atom::changeBasis(per[0]).rows(0,1);
	arma::vec2 holeper2 = Atom::changeBasis(per[1]).rows(0,1);
	
	std::vector<size_t> badatoms;
	size_t size = atoms.size();
	badatoms.reserve(size);
	
	for(size_t a = 0; a < size; a++)
	{
		if(arma::norm(holecent-Atom::changeBasis(atoms[a])(arma::span(0,1)),2) < R && ( st == singlelayer || atoms[a](2) == 0))
			badatoms.push_back(a);
	}

	holecent += holeper1;
	for(size_t a = 0; a < size; a++)
	{
		if(arma::norm(holecent-Atom::changeBasis(atoms[a])(arma::span(0,1)),2) < R && ( st == singlelayer || atoms[a](2) == 0))
			badatoms.push_back(a);
	}

	holecent += holeper2-holeper1;
	for(size_t a = 0; a < size; a++)
	{
		if(arma::norm(holecent-Atom::changeBasis(atoms[a])(arma::span(0,1)),2) < R && ( st == singlelayer || atoms[a](2) == 0))
			badatoms.push_back(a);
	}

	holecent += holeper1;
	for(size_t a = 0; a < size; a++)
	{
		if(arma::norm(holecent-Atom::changeBasis(atoms[a])(arma::span(0,1)),2) < R && ( st == singlelayer || atoms[a](2) == 0))
			badatoms.push_back(a);
	}

	Atom::filter::erase_indices(atoms,badatoms);

	return atoms;
}

std::vector<arma::ivec4> Generator::singletriangulargalrectangle_periodicity(size_t L, double R, const stackingtype &st)
{
	std::vector<arma::ivec4> origper = inftriangulargal_periodicity(L, R, st);
	
	std::vector<arma::ivec4> per;
	
	per.push_back(origper[0]+origper[1]);
	per.push_back(origper[0]-origper[1]);
	
	return per;
}

std::vector<arma::ivec4> Generator::triangulargalrectangle_atoms(size_t L, double R, size_t Nx, size_t Ny, const stackingtype &st)
{
	std::vector<arma::ivec4> singlecell = singletriangulargalrectangle_atoms(L, R, st);
	std::vector<arma::ivec4> singledim = singletriangulargalrectangle_periodicity(L, R, st);

	const size_t size = singlecell.size();

	std::vector<arma::ivec4> result(size*Nx*Ny);
	
	_Cilk_for(size_t i = 0; i < Nx; i++)
	{
		_Cilk_for(size_t j = 0; j < Ny; j++)
		{
			for(size_t a = 0; a < size; a++)
			{
				result[i*Ny*size+j*size+a] = singlecell[a]+singledim[0]*j+singledim[1]*i;
			}
		}
	}

	return result;
}

#include <unordered_map>

class VecHasher
{
public:
	size_t operator()(const arma::uvec4& k) const
    {
      using std::hash;

      // Compute individual hash values for first,
      // second and third and combine them using XOR
      // and bit shifting:

      return ((hash<size_t>()(k(0)))
               ^ (hash<size_t>()(k(1)))
			   ^ (hash<size_t>()(k(2)))
			   ^ (hash<size_t>()(k(3))));
    }
};

class VecEquality
{
public:
	bool operator()(const arma::uvec4& k1, const arma::uvec4& k2) const
    {
      return k1(0) == k2(0) && k1(1) == k2(1) && k1(2) == k2(2) && k1(3) == k2(3);
    }
};

std::vector<arma::ivec4> Generator::triangulargalrectangle_atoms(size_t L, double R, size_t Nx, size_t Ny, const stackingtype &st, size_t padd)
{
	std::unordered_map<arma::uvec4,std::vector<arma::ivec4>,VecHasher,VecEquality> cells;

	std::vector<arma::ivec4> singledim = singletriangulargalrectangle_periodicity(L, R, st);

	size_t size = ((3*L+2*padd)*(2*L+padd*2)*4);

	std::vector<arma::ivec4> result;
	result.reserve(size*Nx*Ny);
	
	arma::ivec4 up("1 -1 0 0");
	arma::ivec4 right("1 1 0 0");
	
	for(size_t i = 0; i < Nx; i++)
	{
		for(size_t j = 0; j < Ny; j++)
		{
			arma::uvec4 paddspec(arma::fill::zeros);
			if(i==0)
				paddspec(0) = padd;
			if(i==Nx-1)
				paddspec(1) = padd;
			if(j==0)
				paddspec(2) = padd;
			if(j==Ny-1)
				paddspec(3) = padd;

			if(cells[paddspec].empty())
				cells[paddspec] = singletriangulargalrectangle_atoms(L, R, st, paddspec);

			size = cells[paddspec].size();

			for(size_t a = 0; a < size; a++)
			{
				result.push_back((cells[paddspec][a]+singledim[0]*j+singledim[1]*i+( (i>0)? padd : 0 )*right+( (j>0)? padd : 0 )*up).eval());
			}
		}
	}

	return result;
}

std::vector<arma::ivec4> Generator::inftriangulargalrectangle_atoms(size_t L, double R, size_t Nx, size_t Ny, const stackingtype &st)
{
	std::unordered_map<arma::uvec4,std::vector<arma::ivec4>,VecHasher,VecEquality> cells;

	std::vector<arma::ivec4> singledim = singletriangulargalrectangle_periodicity(L, R, st);

	size_t size = ((3*L)*(2*L)*4);

	std::vector<arma::ivec4> result;
	result.reserve(size*Nx*Ny);
		
	for(size_t i = 0; i < Nx; i++)
	{
		for(size_t j = 0; j < Ny; j++)
		{
			arma::uvec4 paddspec(arma::fill::zeros);

			if(cells[paddspec].empty())
				cells[paddspec] = singletriangulargalrectangle_atoms(L, R, st, paddspec);

			size = cells[paddspec].size();

			for(size_t a = 0; a < size; a++)
			{
				result.push_back((cells[paddspec][a]+singledim[0]*j+singledim[1]*i).eval());
			}
		}
	}

	return result;
}


std::vector<arma::ivec4> Generator::inftriangulargalrectangle_periodicity(size_t L, double R, size_t Nx, size_t Ny, const stackingtype &st)
{
	std::vector<arma::ivec4> origper = singletriangulargalrectangle_periodicity(L, R, st);
	
	std::vector<arma::ivec4> per(2);
	
	per[0] = origper[0]*Ny;
	per[1] = origper[1]*Nx;

	return per;
}

std::vector<arma::ivec4> Generator::singletriangulargalrectangle_atoms_noholes(size_t L, const stackingtype &st, arma::uvec padd)
{
	// First generate the rectangle of atoms:

	std::vector<arma::ivec4> atoms;
	atoms.reserve((3*L+padd(0)+padd(1))*(2*L+padd(2)+padd(3))*4);

	static const auto atomGen = [] (long i,long j,long h,long l) 
	{
		arma::ivec4 result;
		result(0) = i;
		result(1) = j;
		result(2) = h;
		result(3) = l;
		return result;
	};
	
	for(long l = 0; l < 2; l++)
	{
		for(long j = -(long)L-2*(long)padd(0)-2*(long)padd(1)-2*(long)padd(2)-2*(long)padd(3); j < 8*(long)L; j++)
		{
			for(long i = std::max(-j,j-( (l == 1) ? 1 : 0 )); i < std::min(2*(long)L+j-1+( (l == 1) ? 0 : 1 ) + 2*(long)padd(2) + 2*(long)padd(3),6*(long)L-j + 2*(long)padd(0) + 2*(long)padd(1)); i++)
			{
				for(long h = 0; h < ((st != singlelayer) ? 2: 1); h++)
				{
					atoms.push_back(atomGen(i,j,h,l));
				}
			}
		} 
	}
	
	return atoms;
}

#include <cilk/reducer_list.h>

std::vector<arma::ivec4> Generator::inftriangulargalrectangle_atoms_randomholes(size_t L, double R, size_t Nx, size_t Ny, const stackingtype &st, double randrad)
{
	std::vector<arma::ivec4> singlecell = singletriangulargalrectangle_atoms_noholes(L, st);

	std::vector<arma::ivec4> singledim = singletriangulargalrectangle_periodicity(L, R, st);
	
	size_t size = singlecell.size();

	std::vector<arma::ivec4> result(size*Ny*Nx);

	_Cilk_for(size_t i = 0; i < Nx; i++)
	{
		_Cilk_for(size_t j = 0; j < Ny; j++)
		{
			for(size_t a = 0; a < size; a++)
			{
				result[i*Ny*size+j*size+a] = singlecell[a]+singledim[0]*j+singledim[1]*i;
			}
		}
	}
		
	// Now we make the holes.
	arma::vec2 holecent;
	holecent(0) = -((long)L % 2)*0.5;
	holecent(1) = 0.5/sqrt(3.)+((long)L-1)*0.5*sqrt(3.);
	
	std::vector<arma::ivec4> per = inftriangulargal_periodicity(L, R, st);

	arma::vec2 holeper1 = Atom::changeBasis(per[0]);
	arma::vec2 holeper2 = Atom::changeBasis(per[1]);
	
	std::vector<arma::vec2> holes;
	holes.reserve(Nx*Ny*2+Nx+Ny);

	static const auto randVec2 = [] () {
		arma::vec2 rand = arma::randu<arma::vec>(2);
		arma::vec2 result;
		result(0) = cos(rand(0)*2*arma::datum::pi)*rand(1);
		result(1) = sin(rand(0)*2*arma::datum::pi)*rand(1);
		return result;
	};
	
	for(size_t i = 0; i < Nx; i++)
	{
		for(size_t j = 0; j < Ny; j++)
		{
			// There are two hole for each cell, and edge cell-hole are periodic!
			arma::vec2 randleft = randVec2()*randrad;
			arma::vec2 randtop = randVec2()*randrad;
			
			holes.push_back(holecent+randleft+i*(holeper1+holeper2)+j*(holeper1-holeper2));
			holes.push_back(holecent+randleft+holeper2+i*(holeper1+holeper2)+j*(holeper1-holeper2));
			
			if(i==0)
				holes.push_back(holecent+randleft+Nx*(holeper1+holeper2)+j*(holeper1-holeper2));
			if(j==0)
				holes.push_back(holecent+randleft+holeper2+i*(holeper1+holeper2)+Ny*(holeper1-holeper2));
		}
	}

	size_t atomsize = result.size();
	size_t holesize = holes.size();
	
	std::vector<bool> removeatoms(atomsize,false);

	_Cilk_for(size_t a = 0; a < atomsize; a++)
	{
		for(size_t h = 0; h < holesize && !removeatoms[a]; h++)
		{
			if(arma::norm(holes[h]-Atom(result[a]).xy(),2) < R && ( st == singlelayer || result[a](2) == 0))
				removeatoms[a] = true;
		}
	}
	
	std::vector<size_t> badatoms;
	badatoms.reserve(atomsize);
	
	for(size_t a = 0; a < atomsize; a++)
	{
		if(removeatoms[a])
			badatoms.push_back(a);
	}

	Atom::filter::erase_indices(result,badatoms);

	return result;
}

#include "unit.h"

void Generator::recursive_edge_etching(std::vector<arma::ivec4> &unfilteredatomlist, const std::vector<arma::ivec4> & periodicity, long etchstepstogo)
{
	Atom::list atoms;
	imat per;
	imat offsets;
	
	const size_t origsize = unfilteredatomlist.size();
	const size_t persize = periodicity.size();

	atoms.reserve(origsize);

	for(size_t a = 0; a < origsize; a++)
		atoms.push_back(Atom(unfilteredatomlist[a]));

	per.set_size(4,persize);
	
	for(size_t p = 0; p < persize; p++)
		per.col(p) = periodicity[p];

	offsets = Unit::recursive_find_p_vectors(per);
	
	Atom::reference::grid grid(&atoms,&offsets);

	grid();
	
	const size_t size = atoms.size();
	static const double invsqrtthree = 1./sqrt(3.);
	
	std::vector<size_t> badatoms;
	
	for(size_t a = 0; a < size; a++)
	{		
		int neighbours = 0;
		for(long i = -1; i <= 1; i++)
		{
			for(long j = -1; j <= 1; j++)
			{
				const Atom::reference::list &list = grid.getSurrounding(atoms[a],i,j);
				size_t subsize = list.size();
				for(size_t k = 0; k < subsize; k++)
				{
					// All nearest neighbours are equidistant from the atom.
					double dist = arma::norm((atoms.at(a).xy()-list[k].xy()),2);
					if(atoms[a].h() == list[k].h() && dist > 0.9*invsqrtthree && dist < 1.1*invsqrtthree)
						neighbours++;
				}
			}
		}
		
		if(neighbours == 2 && arma::randu<vec>(1).eval()(0) > 0.1)
			badatoms.push_back(a);
		if(neighbours < 2)
			badatoms.push_back(a);
	}
	
	Atom::filter::erase_indices(unfilteredatomlist,badatoms);
}